"""
Jieun Oh

bitalloc2.py -- Jieun's implementation of bitalloc function for Music 422 HW5

"""
import numpy as np


def BitAlloc(bitBudget, maxMantBits, nBands, nLines, SMR):
    '''Allocates bits to scale factor bands so as to flatten the NMR across the spectrum'''
    '''     Arguments:
               bitBudget is total number of mantissa bits to allocate
               maxMantBits is max mantissa bits that can be allocated per line
               nBands is total number of scale factor bands
               nLines[nBands] is number of lines in each scale factor band
               SMR[nBands] is signal-to-mask ratio in each scale factor band
        
            Return:
                bits[nBands] is number of bits allocated to each scale factor band
        
            Logic:
               Maximizing SMR over blook gives optimization result that:
                   R(i) = P/N + (1 bit/ 6 dB) * (SMR[i] - avgSMR)
               where P is the pool of bits for mantissas and N is number of bands
               This result needs to be adjusted if any R(i) goes below 2 (in which
               case we set R(i)=0) or if any R(i) goes above maxMantBits (in
               which case we set R(i)=maxMantBits).  (Note: 1 Mantissa bit is
               equivalent to 0 mantissa bits when you are using a midtread quantizer.)
               We will not bother to worry about slight variations in bit budget due
               rounding of the above equation to integer values of R(i).
    '''

    bits = np.empty(nBands, dtype=np.int32)             # answer


    P = bitBudget

    Kpassed = 0       # to be set to # bins with positive SMR?? (check this!!!). right now, total # bins
    summation = 0
    for i in range(nBands):
        Kpassed = Kpassed + nLines[i]
        summation = summation + nLines[i] * SMR[i]

    #print("nLines: ", nLines)
    #print("Kpassed: ", Kpassed)
    ##print("SMR: ", SMR)

    '''
    summation = 0
    for subband in range(nBands):
       if (SMR[subband] > 0):
           Kpassed = Kpassed +1
           summation = summation + nLines[subband] * SMR[subband]
    '''

    for subband in range(nBands):
       bitsOpt =( (P/Kpassed)+ 0.05*np.log2(10)*( SMR[subband]-(1/Kpassed)*(summation)) )
       bitsTemp =  np.cast['int'] (bitsOpt)
       if bitsTemp<0:
          bitsTemp = 0
       if bitsTemp>maxMantBits:
          bitsTemp = maxMantBits

       bits[subband] = bitsTemp

       
    #print "bits #1: ", bits
    bitsUsed = sum(nLines*bits)
    #print "used ", bitsUsed, " out of ", P


    # add 1 more bits from lowest to highest bands, excluding the last (25th band), until we run out
    stop = False
    while not stop:
       for subband in range(nBands-1):
          if bitsUsed + nLines[subband] < P:
             bits[subband] = bits[subband]+1
             bitsUsed = bitsUsed + nLines[subband]
          else:
             stop = True

    #print "bits #2: ", bits
    #print "after keep adding ", bitsUsed, " out of ", P


    # final adjustment
    for subband in range(nBands):
       if bits[subband] ==1:
          if bitsUsed < P:
             bits[subband] = 2
             bitsUsed = bitsUsed + nLines[subband]
          else:
             bits[subband] = 0
             bitsUsed = bitsUsed - nLines[subband]

    ##print "bits #final: ", bits
    #print "after taking care of 1-bit assn, used ", bitsUsed, " out of ", P
            
    return bits



# --------------------------Testing--------------------------------#

#if __name__=="__main__":
    
   


    
